package bigdata.hourglass.engine.repl

import java.io.ByteArrayOutputStream
import java.util.{Map => JMap}

import bigdata.hourglass.common.Logging
import bigdata.hourglass.engine.core.SparkRuntime
import org.apache.spark.SparkConf
import org.json4s.DefaultFormats
import org.json4s.JsonAST._
import org.json4s.JsonDSL._

import scala.tools.nsc.interpreter.Results

/**
 * @author leo.jie (weixiao.me@aliyun.com)
 * @organization DataReal
 * @version 1.0
 * @website https://www.jlpyyf.com
 * @date 2019/12/8 8:08 下午
 * @since 1.0
 */
object AbstractSparkInterpreter {
  private[repl] val KEEP_NEWLINE_REGEX = """(?<=\n)""".r
}

abstract class AbstractSparkInterpreter extends Interpreter with Logging {

  import AbstractSparkInterpreter._

  private implicit def formats: DefaultFormats.type = DefaultFormats

  protected val outputStream = new ByteArrayOutputStream()

  protected def isStarted: Boolean

  protected def interpret(code: String): Results.Result

  protected def bind(name: String, tpe: String, value: Object, modifier: List[String]): Unit

  def sparkCreateContext(params: JMap[Any, Any], conf: SparkConf): SparkConf = {
    val sparkRuntime = SparkRuntime.getOrCreate(params)
    val spark = sparkRuntime.createRuntime(conf)
    bind("spark", spark.getClass.getCanonicalName, spark, List("""@transient"""))
    execute("import org.apache.spark.SparkContext._")
    execute("import spark.implicits._")
    execute("import spark.sql")
    execute("import org.apache.spark.sql.functions._")
    spark.sparkContext.getConf
  }

  /** Shut down the interpreter. */
  override def close(): Unit = {}

  /**
   * Execute the code and return the result, it may
   * take some time to execute.
   */
  override def execute(code: String): Interpreter.ExecuteResponse =
    restoreContextClassLoader {
      require(isStarted)
      executeLines(code.trim.split("\n").toList, Interpreter.ExecuteSuccess(JObject(
        (TEXT_PLAIN, JString(""))
      )))
    }

  private class TypesDoNotMatch extends Exception

  private def convertTableType(value: JValue): String = {
    value match {
      case JNothing | JNull => "NULL_TYPE"
      case JBool(_) => "BOOLEAN_TYPE"
      case JString(_) => "STRING_TYPE"
      case JInt(_) => "BIGINT_TYPE"
      case JDouble(_) => "DOUBLE_TYPE"
      case JDecimal(_) => "DECIMAL_TYPE"
      case JArray(arr) =>
        if (allSameType(arr.iterator)) {
          "ARRAY_TYPE"
        } else {
          throw new TypesDoNotMatch
        }
      case JObject(obj) =>
        if (allSameType(obj.iterator.map(_._2))) {
          "MAP_TYPE"
        } else {
          throw new TypesDoNotMatch
        }
      case _ => throw new TypesDoNotMatch
    }
  }

  private def allSameType(values: Iterator[JValue]): Boolean = {
    if (values.hasNext) {
      val type_name = convertTableType(values.next())
      values.forall(value => type_name.equals(convertTableType(value)))
    } else {
      true
    }
  }

  @scala.annotation.tailrec
  private def executeLines(lines: List[String], resultFromLastLine:
  Interpreter.ExecuteResponse): Interpreter.ExecuteResponse = {
    lines match {
      case Nil => resultFromLastLine
      case head :: tail =>
        val result = executeLine(head)

        result match {
          case Interpreter.ExecuteIncomplete() =>
            tail match {
              case Nil =>
                // ExecuteIncomplete could be caused by an actual incomplete statements (e.g. "sc.")
                // or statements with just comments.
                // To distinguish them, reissue the same statement wrapped in { }.
                // If it is an actual incomplete statement, the interpreter will return an error.
                // If it is some comment, the interpreter will return success.
                executeLine(s"{\n$head\n}") match {
                  case Interpreter.ExecuteIncomplete() | Interpreter.ExecuteError(_, _, _) =>
                    // Return the original error so users won't get confusing error message.
                    result
                  case _ => resultFromLastLine
                }
              case next :: nextTail =>
                executeLines(head + "\n" + next :: nextTail, resultFromLastLine)
            }
          case Interpreter.ExecuteError(_, _, _) =>
            result

          case Interpreter.ExecuteAborted(_) =>
            result

          case Interpreter.ExecuteSuccess(e) =>
            val mergedRet = resultFromLastLine match {
              case Interpreter.ExecuteSuccess(s) =>
                // Because of SparkMagic related specific logic, so we will only merge text/plain
                // result. For any magic related output, still follow the old way.
                if (s.values.contains(TEXT_PLAIN) && e.values.contains(TEXT_PLAIN)) {
                  val lastRet = s.values.getOrElse(TEXT_PLAIN, "").asInstanceOf[String]
                  val currRet = e.values.getOrElse(TEXT_PLAIN, "").asInstanceOf[String]
                  if (lastRet.nonEmpty && currRet.nonEmpty) {
                    Interpreter.ExecuteSuccess(TEXT_PLAIN -> s"$lastRet$currRet")
                  } else if (lastRet.nonEmpty) {
                    Interpreter.ExecuteSuccess(TEXT_PLAIN -> lastRet)
                  } else if (currRet.nonEmpty) {
                    Interpreter.ExecuteSuccess(TEXT_PLAIN -> currRet)
                  } else {
                    result
                  }
                } else {
                  result
                }

              case _ => result
            }

            executeLines(tail, mergedRet)
        }
    }
  }

  private def executeLine(code: String): Interpreter.ExecuteResponse = {
    scala.Console.withOut(outputStream) {
      interpret(code) match {
        case Results.Success =>
          Interpreter.ExecuteSuccess(
            TEXT_PLAIN -> readStdout()
          )
        case Results.Incomplete => Interpreter.ExecuteIncomplete()
        case Results.Error =>
          val (eName, traceBack) = parseError(readStdout())
          Interpreter.ExecuteError("Error", eName, traceBack)
      }
    }
  }

  protected[repl] def parseError(stdout: String): (String, Seq[String]) = {
    val lines = KEEP_NEWLINE_REGEX.split(stdout)
    val eName = lines.headOption.map(_.trim).getOrElse("unknown error")
    val traceBack = lines.tail

    (eName, traceBack)
  }

  protected def restoreContextClassLoader[T](fn: => T): T = {
    val currentClassLoader = Thread.currentThread().getContextClassLoader
    try {
      fn
    } finally {
      Thread.currentThread().setContextClassLoader(currentClassLoader)
    }
  }

  private def readStdout(): String = {
    val output = outputStream.toString("UTF-8")
    outputStream.reset()
    output
  }

}


